from typing import re

from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import ChatPromptTemplate
from loguru import logger
from pydantic import Field, BaseModel

from a0base.base_llm import pop_llm


# 定义数据模型
class Answer(BaseModel):
    content: str = Field(description="回答内容")
    tags: str = Field(description="标签，格式为《标签1》《标签2》")


# 自定义解析器
class CustomPydanticOutputParser(PydanticOutputParser):
    def parse(self, text: str) -> Answer:
        # 将数组格式转换为《》格式
        tags_match = re.search(r'"tags"\s*:\s*\[(.*?)\]', text, re.DOTALL)
        if tags_match:
            tags_list = re.findall(r'"([^"]+)"', tags_match.group(1))
            tags_string = "".join([f"《{tag}》" for tag in tags_list])
            text = re.sub(r'"tags"\s*:\s*\[.*?\]', f'"tags": "{tags_string}"', text, flags=re.DOTALL)
        return super().parse(text)


# 创建解析器
parser = CustomPydanticOutputParser(pydantic_object=Answer)

# 获取格式化指令（这会告诉模型如何以 XML 格式输出）
format_instructions = parser.get_format_instructions()

# 创建提示模板
chat_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是AI助手，请按JSON格式简短回答\n{format_instructions}"),
    ("human", "{question}")
])

# 3-使用指定的角色和问题生成具体的提示内容
prompt = chat_prompt.invoke({
    "question": "什么是LangChain",
    "format_instructions": parser.get_format_instructions()
})
# prompt = chat_prompt.format_messages(question="什么是LangChain", format_instructions=format_instructions)
logger.info(prompt)

# 调用模型获取回答结果
response = pop_llm.invoke(prompt)
logger.info(f"解析后的结构化结果:\n{response}")
logger.info(f"结果: {response.content}")
# 打印类型
logger.info(f" 结果response.content类型: {type(response.content)}")
logger.info(f"结果response类型: {type(response)}")
